// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/common/encode/base64.h>
#include <pollux/dwio/common/encryption/encryption.h>
#include <pollux/common/exception/exception.h>

namespace kumo::pollux::dwio::common::encryption::test {
    class TestEncryption {
    public:
        void setKey(const std::string &key) {
            key_ = key;
        }

        const std::string &getKey() const {
            return key_;
        }

        std::unique_ptr<melon::IOBuf> encrypt(melon::StringPiece input) const {
            ++count_;
            auto encoded = pollux::encoding::Base64::encodeUrl(input);
            return melon::IOBuf::copyBuffer(key_ + encoded);
        }

        std::unique_ptr<melon::IOBuf> decrypt(melon::StringPiece input) const {
            ++count_;
            std::string key{input.begin(), key_.size()};
            DWIO_ENSURE_EQ(key_, key);
            auto decoded = pollux::encoding::Base64::decodeUrl(melon::StringPiece{
                input.begin() + key_.size(), input.size() - key_.size()
            });
            return melon::IOBuf::copyBuffer(decoded);
        }

        size_t getCount() const {
            return count_;
        }

    private:
        std::string key_;
        mutable std::atomic<size_t> count_;
    };

    class TestEncrypter : public TestEncryption, public Encrypter {
    public:
        const std::string &getKey() const override {
            return TestEncryption::getKey();
        }

        std::unique_ptr<melon::IOBuf> encrypt(
            melon::StringPiece input) const override {
            return TestEncryption::encrypt(input);
        }

        std::unique_ptr<Encrypter> clone() const override {
            auto encrypter = std::make_unique<TestEncrypter>();
            encrypter->setKey(getKey());
            return encrypter;
        }
    };

    class TestDecrypter : public TestEncryption, public Decrypter {
    public:
        void setKey(const std::string &key) override {
            TestEncryption::setKey(key);
        }

        bool isKeyLoaded() const override {
            return !getKey().empty();
        }

        std::unique_ptr<melon::IOBuf> decrypt(
            melon::StringPiece input) const override {
            return TestEncryption::decrypt(input);
        }

        std::unique_ptr<Decrypter> clone() const override {
            auto decrypter = std::make_unique<TestDecrypter>();
            decrypter->setKey(getKey());
            return decrypter;
        }
    };

    class TestEncryptionProperties : public EncryptionProperties {
    public:
        TestEncryptionProperties(const std::string &key) : key_{key} {
        }

        const std::string &getKey() const {
            return key_;
        }

        size_t hash() const override {
            return std::hash<std::string>{}(key_);
        }

    protected:
        bool equals(const EncryptionProperties &other) const override {
            auto &casted = dynamic_cast<const TestEncryptionProperties &>(other);
            return key_ == casted.key_;
        }

    private:
        std::string key_;
    };

    class TestEncrypterFactory : public EncrypterFactory {
    public:
        std::unique_ptr<Encrypter> create(
            EncryptionProvider provider,
            const EncryptionProperties &props) override {
            DWIO_ENSURE_EQ(provider, EncryptionProvider::Unknown);
            auto encrypter = std::make_unique<TestEncrypter>();
            encrypter->setKey(
                dynamic_cast<const TestEncryptionProperties &>(props).getKey());
            return encrypter;
        }
    };

    class TestDecrypterFactory : public DecrypterFactory {
    public:
        std::unique_ptr<Decrypter> create(EncryptionProvider provider) override {
            DWIO_ENSURE_EQ(provider, EncryptionProvider::Unknown);
            return std::make_unique<TestDecrypter>();
        }
    };
} // namespace kumo::pollux::dwio::common::encryption::test
